package cn.getech.data.metadata.utils;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpEntity;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.HashMap;

/**
 * @author 李海峰
 * @version V1.0
 * @Package cn.getech.data.development.atlas.restapi.util.model
 * @date 2021/7/1 21:44
 */
public class AtlasLineage {
    private static final String BASE_URI = "/api/atlas/v2/search/basic";
    private static final String LINEAGE_URI = "/api/atlas/v2/lineage";
    private static final String HIVE_TABLE = "hive_table";
    private static final String HIVE_PROCESS = "hive_process";
    private static final String HIVE_COLUN = "hive_column";
    private static final String HIVE_COLUN_LINEAGE = "hive_column_lineage";


    private static String atlasServerUrl;   //Atlas ip:port
    private static String userName;        //Atlas 用户名
    private static String password;        //Atlas 密码
    private static String typeName;        //查询类型
    private static String dbName;          //库名
    private static String tableName;       //表名
    private static String tableColumn;     //字段名
    private static int depth;              //展示层级 3,---三层展示， >3---全量展示

    private static Logger log = LoggerFactory.getLogger(AtlasLineage.class);


    /**
     * @param args atlasServerUrl   atlas host:port
     *             username         用户名
     *             passsword        密码
     *             typeName         查询类型（ 库 /表/ 字段/...  hive表是 hive_table / hive_column）
     *             dbName           数据库名（ hive库名 test）
     *             tableName        表名
     *             tableColumn      查询的字段
     *             depth            是否全展示血缘 =3--三层展示， >3 全量展示
     * @return 经过转化的map
     */
    public static void main(String[] args) {

        ArrayList<HashMap<String, String>> res = AtlasLineage.getLineageByTable("http://cdh-hadoop-3:21001",
                "admin", "admin", "hive_column", "test",
                "test_data03", "id", 4);
        log.info(JSONObject.toJSONString(res));
    }

    /**
     * 外部调用接口
     *
     * @param serverUrl
     * @param username
     * @param password
     * @param typeName
     * @param dbName
     * @param tableName
     * @param tableColumn
     * @param depth
     * @return 经过转化的map
     */
    public static ArrayList<HashMap<String, String>> getLineageByTable(String serverUrl, String username, String password,
                                                            String typeName, String dbName,
                                                            String tableName,
                                                            String tableColumn, int depth) {
        atlasServerUrl = serverUrl;
        userName = username;
        AtlasLineage.password = password;
        AtlasLineage.typeName = typeName;
        AtlasLineage.dbName = dbName;
        AtlasLineage.tableName = tableName;
        AtlasLineage.tableColumn = tableColumn;
        AtlasLineage.depth = depth;

        return getLineage();

    }

    /**
     * @return 血缘
     */
    private static ArrayList<HashMap<String, String>> getLineage() {
        //传参校验
        if (StringUtils.isEmpty(atlasServerUrl) || StringUtils.isEmpty(userName) || StringUtils.isEmpty(password) || StringUtils.isEmpty(typeName)
                || StringUtils.isEmpty(dbName) || StringUtils.isEmpty(tableName)) {
            log.error("Atlas param lost! ");
            return new ArrayList<HashMap<String, String>>();
        } else {
            log.info("Already insert atlas param! ");
        }

        if (null == tableColumn) {
            //表血缘
            JSONObject tableEntity = getEntity();
            String tableGuid = tableEntity.getJSONArray("entities").getJSONObject(0).getString("guid");
            System.out.println("tableGuid  ----------- " + tableGuid);
            if (null == tableGuid) {
                log.error("no tableGUID: " + tableName);
                return new ArrayList<HashMap<String, String>>();
            }

            //根据GUID 查询表血缘
            JSONObject tableLeags = getTableByGUID(tableGuid);

            //数据转换 返回
            return tableTrans(tableLeags, depth);
        } else {
            //字段血缘
            JSONObject columnEntity = getEntity();
            String columnGuid = columnEntity.getJSONArray("entities").getJSONObject(0).getString("guid");
            System.out.println("columnGUID  ----------- " + columnGuid);
            if (null == columnGuid) {
                log.error("no columnGUID: " + tableColumn);
                new ArrayList<HashMap<String, String>>();
            }

            //根据GUID 查询字段血缘
            JSONObject columnLeags = getTableByGUID(columnGuid);

            //数据转换 返回
            return columnTrans(columnLeags, depth);

        }
    }


    /**
     * 构造Basic Auth认证头信息
     *
     * @return
     */
    private static String getHeader() {
        String auth = userName + ":" + password;
        byte[] encodedAuth = Base64.encodeBase64(auth.getBytes(Charset.forName("US-ASCII")));
        String authHeader = "Basic " + new String(encodedAuth);
        return authHeader;
    }

    /**
     * 转化为需要的格式返回
     *
     * @param tableLeags
     * @param depth
     * @return table 全量/3层
     */
    private static ArrayList<HashMap<String, String>> tableTrans(JSONObject tableLeags, int depth) {
        ArrayList<HashMap<String, String>> allLayersMap = trans2LineageMap(tableLeags);

        //是三层的话截取 这个表有关的上下游表关系
        if (depth > 3) {
            //全量展示
            return allLayersMap;
        } else {
            //三层展示
            return showTable3Layers(allLayersMap);
        }
    }

    /**
     * @param columnLeags
     * @param depth
     * @return column 全量/3层
     */
    private static ArrayList<HashMap<String, String>> columnTrans(JSONObject columnLeags, int depth) {
        ArrayList<HashMap<String, String>> allLayersMap = trans2LineageMap(columnLeags);

        //是三层的话截取 这个表有关的上下游表关系
        if (depth > 3) {
            //全量展示
            return allLayersMap;
        } else {
            //三层展示
            return showColumn3Layers(allLayersMap);
        }
    }

    private static ArrayList<HashMap<String, String>> showTable3Layers(ArrayList<HashMap<String, String>> allLayersMap) {
        ArrayList<HashMap<String, String>> show3LayersMap = new ArrayList<HashMap<String, String>>();
        allLayersMap.forEach(item -> {
            item.forEach((key, value) -> {
                if (key.equals(dbName + "." + tableName) || value.equals(dbName + "." + tableName)) {
                    HashMap<String, String> newRelastionMap = new HashMap<String, String>() {{
                        put(key, value);
                    }};
                    show3LayersMap.add(newRelastionMap);
                }
            });
        });

        return show3LayersMap;
    }

    private static ArrayList<HashMap<String, String>> showColumn3Layers(ArrayList<HashMap<String, String>> allLayersMap) {
        ArrayList<HashMap<String, String>> show3LayersMap = new ArrayList<HashMap<String, String>>();
        allLayersMap.forEach(item -> {
            item.forEach((key, value) -> {
                if (key.equals(dbName + "." + tableName + "." + tableColumn) || value.equals(dbName + "." + tableName + "." + tableColumn)) {
                    HashMap<String, String> newRelastionMap = new HashMap<String, String>() {{
                        put(key, value);
                    }};
                    show3LayersMap.add(newRelastionMap);
                }
            });
        });

        return show3LayersMap;
    }

    private static ArrayList<HashMap<String, String>> trans2LineageMap(JSONObject tableLeags) {
        ArrayList<HashMap<String, String>> allRelationsMap = new ArrayList<>();
        JSONArray relations = tableLeags.getJSONArray("relations");

        if ((null == relations) || (relations.size() < 1)) {
            return allRelationsMap;
        } else {
            String fromProcess = null;
            String toProcess = null;
            if (null != tableColumn) {
                fromProcess = HIVE_COLUN;
                toProcess = HIVE_COLUN_LINEAGE;

            } else {
                fromProcess = HIVE_TABLE;
                toProcess = HIVE_PROCESS;
            }

            for (int i = 0; i < relations.size(); i++) {
                JSONObject jsonObject = relations.getJSONObject(i);
                String fromEntityId = jsonObject.getString("fromEntityId");
                String toEntityId = jsonObject.getString("toEntityId");

                if (tableLeags.getJSONObject("guidEntityMap").getJSONObject(fromEntityId).getString("typeName").equals(fromProcess)) {
                    if (tableLeags.getJSONObject("guidEntityMap").getJSONObject(toEntityId).getString("typeName").equals(toProcess)) {

                        for (int j = 0; j < relations.size(); j++) {
                            JSONObject jsonObject1 = relations.getJSONObject(j);
                            String fromEntityId1 = jsonObject1.getString("fromEntityId");
                            String toEntityId1 = jsonObject1.getString("toEntityId");
                            if (tableLeags.getJSONObject("guidEntityMap").getJSONObject(fromEntityId1).getString("typeName").equals(toProcess) && fromEntityId1.equals(toEntityId)) {
                                if (tableLeags.getJSONObject("guidEntityMap").getJSONObject(toEntityId1).getString("typeName").equals(fromProcess)) {
                                    String fromName = tableLeags.getJSONObject("guidEntityMap").getJSONObject(fromEntityId).getJSONObject("attributes").getString("qualifiedName").split("@")[0];
                                    String toName = tableLeags.getJSONObject("guidEntityMap").getJSONObject(toEntityId1).getJSONObject("attributes").getString("qualifiedName").split("@")[0];
                                    HashMap<String, String> newRelastionMap = new HashMap<String, String>() {{
                                        put(fromName, toName);
                                    }};
                                    if (!allRelationsMap.contains(newRelastionMap)) {
                                        allRelationsMap.add(newRelastionMap);
                                    }
                                }
                            }
                        }
                    }
                }
            }
            return allRelationsMap;
        }
    }


    /**
     * @return 表实体
     */
    private static JSONObject getEntity() {
        CloseableHttpClient client = HttpClients.createDefault();
        HttpPost post = new HttpPost(atlasServerUrl + BASE_URI);
        //请求体消息
        String pushObject = "{\"typeName\":\"" +
                typeName +
                "\",\"excludeDeletedEntities\":true,\"includeClassificationAttributes\":true,\"includeSubTypes\":true,\"includeSubClassifications\":true,\"limit\":25,\"offset\":0,\"entityFilters\":{\"condition\":\"AND\",\"criterion\":[{\"attributeName\":\"name\",\"operator\":\"=\",\"attributeValue\":\"" +
                "atlas_test_test1" +
                "\"},{\"attributeName\":\"qualifiedName\",\"operator\":\"=\",\"attributeValue\":\"" +
                dbName +
                "." +
                tableName +
                "@primary\"}]},\"attributes\":[\"qualifiedName\"]}";
        if (null != tableColumn) {
            pushObject = "{\"typeName\":\"" +
                    typeName +
                    "\",\"excludeDeletedEntities\":true,\"includeClassificationAttributes\":true,\"includeSubTypes\":true,\"includeSubClassifications\":true,\"limit\":25,\"offset\":0,\"entityFilters\":{\"condition\":\"AND\",\"criterion\":[{\"attributeName\":\"name\",\"operator\":\"=\",\"attributeValue\":\"" +
                    tableColumn +
                    "\"},{\"attributeName\":\"qualifiedName\",\"operator\":\"=\",\"attributeValue\":\"" +
                    dbName +
                    "." +
                    tableName +
                    "." +
                    tableColumn +
                    "@primary\"}]},\"attributes\":[\"qualifiedName\"]}";
        }
        StringEntity myEntity = new StringEntity(pushObject, ContentType.APPLICATION_JSON); // 构造请求数据
        post.addHeader("Authorization", getHeader()); //可能还放别的参数
        post.setEntity(myEntity); // 设置请求体
        String responseContent = null; // 响应内容
        CloseableHttpResponse response = null;
        try {
            response = client.execute(post);
            if (response.getStatusLine().getStatusCode() == 200) {
                HttpEntity entity = response.getEntity();
                responseContent = EntityUtils.toString(entity, "UTF-8");
                return JSON.parseObject(responseContent);
            }
        } catch (IOException e) {
            log.error("获取血缘错误",e);
        }finally {
            if (response != null) {
                try {
                    response.close();
                } catch (IOException ignored) {
                }
            }
            if (client != null) {
                try {
                    client.close();
                } catch (IOException ignored) {
                }
            }
        }
       return null;

    }

    /**
     * @param guid
     * @return 血缘
     */
    private static JSONObject getTableByGUID(String guid) {
        CloseableHttpClient client = HttpClients.createDefault();
        HttpGet httpGet = new HttpGet(atlasServerUrl + LINEAGE_URI + "/" + guid + "?&3&both");
        httpGet.addHeader("Authorization", getHeader()); //可能还放别的参数
        String responseContent = null; // 响应内容
        CloseableHttpResponse response = null;
        try {
             response = client.execute(httpGet);
            if (response.getStatusLine().getStatusCode() == 200) {
                HttpEntity entity = response.getEntity();
                responseContent = EntityUtils.toString(entity, "UTF-8");
                return JSON.parseObject(responseContent);
            }
        } catch (IOException e) {
            log.error("获取血缘错误",e);
        }finally {
            if (response != null) {
                try {
                    response.close();
                } catch (IOException ignored) {
                }
            }
            if (client != null) {
                try {
                    client.close();
                } catch (IOException ignored) {
                }
            }
        }
      return null;
    }

}
